import typing as tp
import pandas as pd
import numpy as np
from numba import njit


def getWeights(d: float, size: int):
    """迭代法计算差分的权重序列
    
    :param d: 阶数
    :param size: 计算的长度
    """
    # thres>0 drops insignificant weights
    w = [1.]
    for k in range(1, size):
        w_ = -w[-1]/k*(d-k+1)
        w.append(w_)
        w = np.array(w[::-1]).reshape(-1, 1)
    return w


def fracDiff(series: pd.DataFrame,
             d: float,
             thres: float = .01):
    """STANDARD FRACDIFF (EXPANDING WINDOW)
    
    expanding window法计算小数阶差分
    Note 1: For thres=1, nothing is skipped.
    Note 2: d can be any positive fractional, not necessarily bounded [0,1]
    
    :param series: 特征时间序列
    :param d: 阶数
    :param thers: 丢弃不显著权重的阈值
    """
    #1) 按最长的长度计算权重序列
    w = getWeights(d, series.shape[0])
    #2) Determine initial calcs to be skipped based on weight-loss threshold
    w_ = np.cumsum(np.abs(w))
    w_ /= w_[-1]
    skip = w_[w_ > thres].shape[0]
    #3) Apply weights to values
    df = {}
    for name in series.columns:
        seriesF, df_ = series[[name]].fillna(
            method='ffill').dropna(), pd.Series()
        for iloc in range(skip, seriesF.shape[0]):
            loc = seriesF.index[iloc]
            if not np.isfinite(series.loc[loc, name]):  # type: ignore
                continue  # exclude NAs
            df_[loc] = np.dot(w[-(iloc+1):, :].T, seriesF.loc[:loc])[0, 0]  # type: ignore
        df[name] = df_.copy(deep=True)
    df = pd.concat(df, axis=1)
    return df


def getWeights_FFD(d: float, thres: float = 1e-5):
    """迭代法计算差分的权重序列（fixed-window）
    
    :param d: 阶数
    :param thres: 丢弃不显著阶数的阈值
    """
    # thres>0 drops insignificant weights
    w = [1.]
    k = 1
    while True:
        w_ = -w[-1]/k*(d-k+1)
        if abs(w_) > thres:
            w.append(w_)
        else:
            break
    w = np.array(w[::-1]).reshape(-1, 1)
    return w


def fracDiff_FFD(series: pd.DataFrame, 
                 d: float, 
                 thres: float = 1e-5):
    """Constant width window (new solution)
    
    Note 1: thres determines the cut-off weight for the window
    Note 2: d can be any positive fractional, not necessarily bounded [0,1]
    
    :param series: 特征时间序列
    :param d: 阶数
    :param thers: 丢弃不显著权重的阈值
    """
    #1) Compute weights for the longest series
    w = getWeights_FFD(d, thres)
    width = len(w)-1
    #2) Apply weights to values
    df = {}
    for name in series.columns:
        seriesF, df_ = series[[name]].fillna(
            method='ffill').dropna(), pd.Series()
        # todo: 循环速度太慢，需要改进
        for iloc1 in range(width, seriesF.shape[0]):
            loc0, loc1 = seriesF.index[iloc1-width], seriesF.index[iloc1]
            if not np.isfinite(series.loc[loc1, name]):  # type: ignore
                continue  # exclude NAs
            df_[loc1] = np.dot(w.T, seriesF.loc[loc0:loc1])[0, 0]  # type: ignore
        df[name] = df_.copy(deep=True)
    df = pd.concat(df, axis=1)
    return df


@njit
def _ffd_inner_loop(series, width, w):
    res = np.zeros_like(series)
    for iloc1 in range(width, series.shape[0]):
        iloc0 = iloc1 - width
        if not np.isfinite(series[iloc1]):
            continue
        res[iloc1] = np.dot(w.T, series[iloc0:iloc1+1].reshape(-1, 1))[0, 0]
    res[:width] = np.nan
    return res


def fracDiff_FFD_col(series: pd.Series, d, thres=1e-5):
    w = getWeights_FFD(d, thres)
    width = len(w) - 1
    seriesF = series.fillna(method='ffill').dropna()
    df_ = _ffd_inner_loop(seriesF.values, width, w)
    df_ = pd.Series(df_.ravel(), index=seriesF.index, name=series.name)
    return df_
